import numpy as np
import time
import pickle
import math
import os
from scipy import interpolate
Defaultdr=0.05
def r_min(rs,gs,scale=10,s=1):
    x=np.linspace(0,rs[-1],rs.shape[-1]*scale)
    soft=interpolate.UnivariateSpline(rs,gs,s=s)
    y=soft(x)
    start=y.argmax()
    return x[np.where(y[start+1:]-y[start:-1]>0)[0][0]+start]

class structure():
    def __init__(self,**kwag):
        if 'cell_parameter' in kwag:
            self.cell_parameter=kwag['cell_parameter']
        else:
            self.cell_parameter=[]
        if 'atmnamelist' in kwag:
            self.atmnamelist=kwag['atmnamelist']
        else:
            self.atmnamelist=[]
        if 'atmnumlist' in kwag:
            self.atmnumlist=kwag['atmnumlist']
        else:
            self.atmnumlist=[]
        if 'frames' in kwag:
            self.frames=kwag['frames']
        else:
            self.frames=[]
        if 'dis' in kwag:
            self.dis=kwag['dis']
        else:
            self.dis=False
    def readXDATCAR(self,filename):
        fp=open(filename)
        line=fp.readline()
        line=fp.readline()
        scale=float(line.strip())
        cell_parameter=np.empty((3,3),dtype=float)
        #读取cell_parameter
        for i in range(0,3):
            line=fp.readline()
            j=0
            for m in line.strip().split():
                cell_parameter[i,j]=float(m)
                j+=1
        self.cell_parameter=cell_parameter*scale
        line=fp.readline()
        self.atmnamelist=line.strip().split()
        line=fp.readline()
        self.atmnumlist=np.array([int(x) for x in line.strip().split()],dtype=int)
        #读取每一帧的位置
        #只能读Direct坐标系
        AtomNum=self.atmnumlist.sum()
        while True:
            frame=np.empty((AtomNum,3),dtype=float)
            try:
                line=fp.readline()
            except:
                break
            if line=="":
                fp.close()
                break
            #StartTime=time.time()
            for i in range(0,AtomNum):
                line=fp.readline()
                j=0
                for m in line.strip().split():
                    num=float(m)
                    frame[i,j]=num-math.floor(num)
                    j+=1
            self.frames.append(frame)
            #EndTime=time.time()
            #print("Reading a frame cost %s"%(EndTime-StartTime))
    def normed(self):
        try:
            return self.is_normed
        except AttributeError:
            i,j=np.nonzero(self.cell_parameter)
            is_normed=np.all(i==j)
            self.is_normed=is_normed
            return is_normed   
    def get_absolutepos(self):
        try:
            return self.absolutepos
        except AttributeError:
            self.absolutepos=np.empty((len(self.frames),self.atmnumlist.sum(),3))
            for f in range(0,len(self.frames)):
                #StartTime=time.time()
                self.absolutepos[f,:,:]=self.frames[f].dot(self.cell_parameter)
                #EndTime=time.time()
                #print("Calculating the absolutepos of a frame cost %s"%(EndTime-StartTime))
            return self.absolutepos
    def min_diss(self,A,B,cp=False,normed=-1):
        if type(cp)==bool and cp==False:
            cp=self.cell_parameter
        if normed==-1:
            normed=self.normed()
        if normed:
            dis_v=np.abs(A-B)
            dis_v=np.where(dis_v<0.5,dis_v,1-dis_v)
            return np.linalg.norm(dis_v.dot(cp))
        diss=np.empty(27,dtype=float)
        di=0
        for mx in (-1,0,1):
            for my in (-1,0,1):
                for mz in (-1,0,1):
                    dis_v=(A-(B+np.array([mx,my,mz]))).dot(cp)
                    diss[di]=np.linalg.norm(dis_v)
                    if diss[di]==0:
                        return 0.0
                    di+=1
        return diss.min()
    def _distance_normed(self):
        cp=self.cell_parameter
        AtomNum=self.atmnumlist.sum()
        FrameNum=len(self.frames)
        self.dis=np.zeros((FrameNum,AtomNum,AtomNum),dtype=float)
        for f in range(0,FrameNum):
            AMatrix=self.frames[f]
            for j in range(1,AtomNum):
                dis_v=abs(AMatrix[j:,:]-AMatrix[:AtomNum-j,:])
                dis_v=np.where(dis_v<0.5,dis_v,1-dis_v)
                dis_diag=np.linalg.norm(dis_v.dot(cp),axis=1)
                np.fill_diagonal(self.dis[f,j:,0:AtomNum-j],dis_diag)
            self.dis[f,:,:]+=self.dis[f,:,:].T
    def _distance_any(self):
        cp=self.cell_parameter
        AtomNum=self.atmnumlist.sum()
        FrameNum=len(self.frames)
        self.dis=np.zeros((FrameNum,AtomNum,AtomNum),dtype=float)
        AabPos=self.get_absolutepos()
        shiftpos=np.empty((27,3),dtype=float)
        shiftindex=0
        for mx in (-1,0,1):
            for my in (-1,0,1):
                for mz in (-1,0,1):
                    shiftpos[shiftindex,:]=np.array([mx,my,mz]).dot(cp)
                    shiftindex+=1
        for f in range(0,FrameNum):
            AabMatrix=AabPos[f,:,:]
            for j in range(1,AtomNum):
                dis_value=np.empty((27,AtomNum-j),dtype=float)
                dis_ab=AabMatrix[j:,:]-AabMatrix[:AtomNum-j,:]
                for m in range(0,27):
                    dis_value[m,:]=np.linalg.norm(dis_ab-shiftpos[m,:],axis=1)
                dis_diag=dis_value.min(axis=0)
                np.fill_diagonal(self.dis[f,j:,0:AtomNum-j],dis_diag)
            self.dis[f,:,:]+=self.dis[f,:,:].T                        
    def distances(self):
        normed=self.normed()
        print("start calculating distance, this will take a lot of time...")
        if normed:
            self._distance_normed()
        else:
            self._distance_any()
    def get_cellV(self):
        try:
            return self.cellV
        except AttributeError:
            self.cellV=abs(np.linalg.det(self.cell_parameter))
            return self.cellV
    def _atm_range(self,I):
        Start=self.atmnumlist[0:I].sum()
        End=self.atmnumlist[0:I+1].sum()
        return Start,End
    def get_rdf(self,alphaI,betaI,dr=False):
        try:
            rdfdict=self.rdf
        except AttributeError:
            self.rdf=dict()
        try:
            Ndict=self.N
        except AttributeError:
            self.N=dict()
        if dr == False:
            for rdf in self.rdf:
                if rdf[0:2]==(alphaI,betaI) or rdf[0:2]==(betaI,alphaI):
                    return self.rdf[rdf]
            dr=Defaultdr
        if (alphaI,betaI,dr) in self.rdf or (betaI,alphaI,dr) in self.rdf:
            if (alphaI,betaI,dr) in self.rdf:
                return self.rdf[(alphaI,betaI,dr)]
            else:
                return self.rdf[(betaI,alphaI,dr)]
        else:
            print("calculating rdf of %s and %s, dr=%s"%(self.atmnamelist[alphaI],self.atmnamelist[betaI],dr))    
        V=self.get_cellV()
        rho=self.atmnumlist[betaI]/V
        alphaStart,alphaEnd=self._atm_range(alphaI)
        betaStart,betaEnd=self._atm_range(betaI)
        #StartTime=time.time()
        RelatedDis=self.dis[:,alphaStart:alphaEnd,betaStart:betaEnd].copy().flatten()
        #EndTime=time.time()
        #print("copy method cost: %s s"%(EndTime-StartTime))
        #StartTime=time.time()
        RelatedDis.sort(kind="mergesort")
        #EndTime=time.time()
        #print("sorting cost: %s s"%(EndTime-StartTime))
        maxr=RelatedDis.max()
        r_bin=np.array(range(0,math.floor(maxr/dr)+1))*dr
        N,edge=np.histogram(RelatedDis,r_bin)
        Frames=len(self.frames)
        g=N/(4*math.pi*r_bin[1:]*r_bin[1:]*dr*rho*Frames*self.atmnumlist[alphaI])
        self.rdf[(alphaI,betaI,dr)]=[r_bin[1:],g]
        self.N[(alphaI,betaI,dr)]=[r_bin[1:],N]
        return self.rdf[(alphaI,betaI,dr)]
    def someDis(self,alphaI,betaI):
        alphaStart,alphaEnd=self._atm_range(alphaI)
        betaStart,betaEnd=self._atm_range(betaI)
        return self.dis[:,alphaStart:alphaEnd,betaStart:betaEnd]
    def getN(self,alphaI,betaI,dr=False):
        try:
            self.N
        except AttributeError:
            self.N=dict()
        try:
            self.rdf
        except AttributeError:
            self.rdf=dict()
        if dr == False:
            for n in self.N:
                if n[0:2]==(alphaI,betaI) or n[0:2]==(betaI,alphaI):
                    return self.N[n]
            dr=Defaultdr
        if (alphaI,betaI,dr) in self.N or (betaI,alphaI,dr) in self.N:
            if (alphaI,betaI,dr) in self.N:
                return self.N[(alphaI,betaI,dr)]
            else:
                return self.N[(betaI,alphaI,dr)]
        else:
            self.get_rdf(alphaI,betaI,dr)
            return self.N[(alphaI,betaI,dr)]        
    def averageCN(self,alphaI,betaI,r_cut=False):
        FrameNum=len(self.frames)
        if r_cut==False:
            rdf=self.get_rdf(alphaI,betaI)
            r_cut=r_min(*rdf)
        N=self.getN(alphaI,betaI)
        Index=np.searchsorted(N[0],r_cut,side='right')
        return N[1][:Index].sum()/FrameNum/self.atmnumlist[alphaI]
    def averageBondLength(self,alphaI,betaI,r_cut=False):
        if r_cut==False:
            rdf=self.get_rdf(alphaI,betaI)
            r_cut=r_min(*rdf)
        rdf=self.get_rdf(alphaI,betaI)
        r=rdf[0]
        g=rdf[1]
        dr=r[2]-r[1]
        Index=np.searchsorted(r,r_cut,side='left')
        r_new=r[:Index+1]
        g_new=g[:Index+1]
        return ((r_new-dr/2)*g_new).sum()/g_new.sum()
    def CN(self,alphaI,betaI,r_cut=False):
        FrameNum=len(self.frames)
        if r_cut==False:
            rdf=self.get_rdf(alphaI,betaI)
            r_cut=r_min(*rdf)
        ABDis=self.someDis(alphaI,betaI).view()
        isBond=np.where(ABDis<r_cut,1,0)
        NumAlpha=self.atmnumlist[alphaI]
        Numbeta=self.atmnumlist[betaI]
        return isBond.reshape((FrameNum,NumAlpha*Numbeta)).sum(axis=1)/NumAlpha
    def BondLength(self,alphaI,betaI,r_cut=False):
        FrameNum=len(self.frames)
        if r_cut==False:
            rdf=self.get_rdf(alphaI,betaI)
            r_cut=r_min(*rdf)
        ABDis=self.someDis(alphaI,betaI).view()
        NumAlpha=self.atmnumlist[alphaI]
        Numbeta=self.atmnumlist[betaI]
        isBond=np.where(ABDis<r_cut,1,0)
        Bondlength=np.where(isBond==1,ABDis,0.0)
        NBond=isBond.reshape((FrameNum,NumAlpha*Numbeta)).sum(axis=1)
        return Bondlength.reshape((FrameNum,NumAlpha*Numbeta)).sum(axis=1)/NBond
        
class XDAT():
    def __init__(self,XDATName,**kwargs):
        self.name=XDATName
        if 'cachename' not in kwargs:
            cachename=".structure."+XDATName
        else:
            cachename=kwargs['cachename']
        self.cachename=cachename
        if 'time' in kwargs:
            timecheck=kwargs['time']
        else:
            timecheck=False
        if timecheck:
            StartTime=time.time()
        if os.path.isfile(cachename) and os.access(cachename, os.R_OK):
            print("Finding cachefile, reading data from cachefile directly……")
            print("Be aware that, cachefile is not always work well, if any error happened at this step, delete the cachefile:%s"%cachename)
            cachefp=open(cachename,'rb')
            newstrc=pickle.load(cachefp)
            cachefp.close()
        else:
            print("No cachefile exist, reading data from xdatcar file……")
            newstrc=structure()
            newstrc.readXDATCAR(XDATName)
            newstrc.distances()
            sfp=open(cachename,'wb')
            pickle.dump(newstrc,sfp)
            sfp.close()
        self.structure=newstrc
        if timecheck:
            EndTime=time.time()
            print("pre-process cost: %s s"%(EndTime-StartTime))
        try:
            self.structure.rdf
            self.rdf_calced=True
        except AttributeError:
            self.rdf_calced=False
    def save(self,silence=False):
        if not silence:
            print("Start saving data...")
        sfp=open(self.cachename,'wb')
        pickle.dump(self.structure,sfp)
        sfp.close()
        if not silence:
            print("Finished saving data...")
    def SomeDistance(self,alphaI,betaI,copy=True):
        if copy:
            return self.structure.someDis(alphaI,betaI).copy()
        else:
            return self.structure.someDis(alphaI,betaI)
    def all_rdf(self):
        for i in range(0,len(self.structure.atmnamelist)):
            for j in range(0,len(self.structure.atmnamelist)):
                if not i==j:
                    self.structure.get_rdf(i,j)
        self.save(silence=True)
        self.rdf_calced=True
def testjob():
    import matplotlib.pyplot as plt
    fname="XDATCAR2"
    xdat=XDAT(fname,time=True)
    #print(xdat.structure.dis[0,3,6])
    #print(xdat.structure.dis[0,8,95])
    #print(xdat.structure.dis[89999,8,95])
    #print(xdat.structure.dis[89999,95,8])
    #print(xdat.structure.dis[89999,5,9])
    if not xdat.rdf_calced:
        xdat.all_rdf()
    SiO=xdat.structure.get_rdf(0,5)
    rcut=r_min(*SiO)
    print(rcut)
    print(xdat.structure.averageCN(0,5,2.35))
    #plt.plot(SiO[0],SiO[1])
    OSi=xdat.structure.get_rdf(5,0)
    rcut=r_min(*OSi)
    print(rcut)
    print(xdat.structure.averageCN(5,0,2.35))
    #plt.plot(OSi[0],OSi[1])
    #plt.show()
    #print(xdat.structure.averageCN(0,5))
    #print(xdat.structure.averageBondLength(0,5))
    #print(xdat.structure.averageBondLength(5,6))
    #CN=xdat.structure.CN(0,5)
    #plt.plot(CN)
    #BondLength=xdat.structure.BondLength(0,5)
    #plt.plot(BondLength)
    #plt.show()
    #print(CN.shape)
if __name__=='__main__':
    testjob()
else:
    print("你正在使用 结构派 后处理脚本")
    print("https://gitee.com/xczics/structure-py")
